#Notes for Graeme:
# 1.  Took your suggestion and made all matrices 7 x 7
# 2. ad1, ad2, ad3 in sampler function now represent vectors representing rows with adults, e.g., 3:5 in the case of os
# 3. Sampler now incorporates situations where the actual population size is below pop.size
# 4. cell.update spits out a list with [[1]] being the population matrices for each species (as a 3D array: rows, columns, species), and [[2]] being the fire status of that cell
# 5. Populations are tracked as a 5D array, of dimensions [7,1,3,x,y], being rows, columns, species, x, and y location.
# 6. Fire history is tracked as a simple matrix, dimensions [x,y] with zeros or ones to record previous fire history for calculating the conditional probability of a burn
# 7. Pop and fire arrays are set up by the initialiser functions, space.init, and fire.space.init, other functions are modified to work off this framework.
# 8. cl.size, and sp.size are variables to set the size of the space, and the size of the cells within the space. Setting them as equal equates to a space with a single grid cell and yields the same answer as the previous code (corrected).  Try it out.
# 9. setting sp.size to a multiiple of cl.size will create a grid and populations will run in each of the grid cells.  Remember that the number of populations increases with the square of the multiple you choose, so it is easy to have a shitload of populations running, and to bog your computer down.

### Functions ###

# Takes population matrices and vector of adult indexes.
# samples up to a maximum number of adults (introduces competition)
sampler<-function(mat1,mat2, mat3, ad1, ad2, ad3, pop.size, comp=TRUE){
	if (comp==FALSE) {
    mat1[ad1]<-ifelse(mat1[ad1]>pop.size, pop.size, mat1[ad1])
    mat2[ad2]<-ifelse(mat2[ad2]>pop.size, pop.size, mat2[ad2])
    mat3[ad3]<-ifelse(mat3[ad3]>pop.size, pop.size, mat3[ad3])
    return(array(c(mat1, mat2, mat3), dim=c(7,1,3)))
	}  
  prob<-c(mat1[ad1], mat2[ad2], mat3[ad3])
	if (sum(prob)==0) {
		out2<-array(c(mat1, mat2, mat3), dim=c(7,1,3))
		return(out2)
	}
	if (sum(prob)<pop.size) pop.size<-sum(prob)
	prob<-prob/sum(prob)
	out<-rmultinom(1, pop.size, prob)
	mat1[ad1]<-out[1:length(ad1)]
	mat2[ad2]<-out[(length(ad1)+1):length(c(ad1,ad2))]
	mat3[ad3]<-out[(length(c(ad1,ad2))+1):length(c(ad1,ad2,ad3))]
	out2<-array(c(mat1, mat2, mat3), dim=c(7,1,3))
	out2		
}

# given a set of populations in a cell and a fire frequency, this calculates (with competition) the next generation
# returns new set of cell populations and the cell fire state
cell.update<-function(pop.or, pop.os, pop.fr, lh.or, lh.os, lh.fr, fire.grid.cell, comp=TRUE){
	if (sum(pop.os[2:5])==0) lh.os[2,1]<-0.025 #quasi density dependent germination
	if (sum(pop.fr[3:7])==0) lh.fr[3,1]<-0.023
	nf.or<-lh.or%*%pop.or
	nf.os<-lh.os%*%pop.os
	nf.fr<-lh.fr%*%pop.fr
	if(fire.grid.cell==T) {
		nf.or[2,1]<-0.55*sum(nf.or[4:5,1])
		nf.or[c(1, 3:5),1]<-0
		nf.os[2:5,1]<-0
		nf.fr[2,1]<-0.85*sum(nf.fr[4:7,1])
		nf.fr[3:7,1]<-0
	}
	out<-sampler(nf.or, nf.os, nf.fr, 3:5, 2:5, 3:7, K, comp)
	out
}

# Initialises the space
# Sets up an array containing the population for each species (a 3D array) with x and y tags
# this makes it a 5D array, which is hard to think about, but should be easy to use... maybe.
space.init<-function(space.size, cell.size, init.or, init.os, init.fr){
	lth<-space.size/cell.size
	out<-array(c(init.or, init.os, init.fr), dim=c(7, 1, 3, lth, lth))
	out
}

# Initialises the matrix of fire events in space
space.fire.init<-function(space.size, cell.size){
	lth<-space.size/cell.size
	out<-array(0, dim=c(lth, lth))	
}

# For each cell in toroidal space, collects eight neighbours plus target cell 
# returns a 4D array: a matrix representing nine rows of x,y coordinates for each cell, so dimensions [9,2,x,y]
neighbours.init<-function(space.size, cell.size){
	lth<-space.size/cell.size
	out<-array(0, dim=c(9,2,lth,lth))
	for (x in 1:lth){
		for (y in 1:lth){
			ifelse((x-1)==0, X<-lth, X<-(x-1)) #correct for indices going to zero
			ifelse((y-1)==0, Y<-lth, Y<-(y-1))
			ifelse((y+1)>lth, Yp<-1, Yp<-(y+1)) # correct for indices going to > space
			ifelse((x+1)>lth, Xp<-1, Xp<-(x+1))
			# Assign neighbours
			out[1,,x,y]<-c(x, y) #record target cell
			out[2,,x,y]<-c(Xp, y)
			out[3,,x,y]<-c(X, y)
			out[4,,x,y]<-c(x, Y)
			out[5,,x,y]<-c(Xp, Y)	
			out[6,,x,y]<-c(X, Y)
			out[7,,x,y]<-c(x, Yp)
			out[8,,x,y]<-c(Xp, Yp)
			out[9,,x,y]<-c(X, Yp)
		}	
	}
	out	#return the array
}

# Dispersal for resprouters
# Note, this function is never called as it has been combined into comb.disp
or.disperse<-function(pop.grid, neighbour.grid, p.disp){
	#get populations of stolons - Graeme, pls check that these are in the second row
	pop.size.grid<-pop.grid[2,1,1,,]	
	if (length(pop.size.grid)==1){
		print("Warning: Space consists of one grid cell - dispersal is meaningless")
		return(pop.grid)	
	}
	# work out number of dispersers
	n.disp.grid<-floor(p.disp*pop.size.grid)
	# disperse the fuckers
	lth<-length(pop.size.grid[,1])
	for (x in 1:lth){
		for (y in 1:lth){
			if (n.disp.grid[x, y]==0) next
			rows<-sample(1:9, n.disp.grid[x,y], replace=T)
			rows<-c(rows, 1:9)
			rows<-table(rows)-1
			pop.grid[2,1,1,x,y]<-pop.grid[2,1,1,x,y]-sum(rows)
			for (i in 1:9){
				pop.grid[2,1,1,neighbour.grid[i,1,x,y],neighbour.grid[i,2,x,y]]<-rows[i]+pop.grid[2,1,1,neighbour.grid[i,1,x,y],neighbour.grid[i,2,x,y]]
			}
		}	
	}
	pop.grid
}


# Dispersal for seeders
# Note, this function is never called as it has been combined into comb.disp
others.disperse<-function(pop.grid, p.disp){
	lth<-length(pop.grid[1,1,1,1,])
	if (lth==1){
		print("Warning: Space consists of one grid cell - dispersal is meaningless")
		return(pop.grid)	
	}
	minus.os<-array(0, dim=rep(lth, 2))
	minus.fr<-array(0, dim=rep(lth, 2))
	pl.os<-0
	pl.fr<-0
	for(x in 1:lth){
		for(y in 1:lth){
			minus.os[x,y]<-floor(p.disp*pop.grid[1,1,2,x,y])
			pl.os<-pl.os+minus.os[x,y]
			minus.fr[x,y]<-floor(p.disp*pop.grid[1,1,3,x,y])
			pl.fr<-pl.fr+minus.fr[x,y]	
		}	
	}
	rain.os.x<-c(sample(1:lth, pl.os, replace=T), rep(1:lth, lth))
	rain.os.y<-c(sample(1:lth, pl.os, replace=T), rep(1:lth, each=lth))
	rain.fr.x<-c(sample(1:lth, pl.fr, replace=T), rep(1:lth, lth))
	rain.fr.y<-c(sample(1:lth, pl.fr, replace=T), rep(1:lth, each=lth))
	rain.os<-table(rain.os.x, rain.os.y)-1
	rain.fr<-table(rain.fr.x, rain.fr.y)-1
	pop.grid[1,1,2,,]<-pop.grid[1,1,2,,]-minus.os+rain.os
	pop.grid[1,1,3,,]<-pop.grid[1,1,3,,]-minus.fr+rain.fr
	pop.grid	
}

# The combined dispersal function.  Will likely look like complete gibberish unless you understand the above two!
comb.disp<-function(pop.grid, neighbour.grid, p.disp.or, p.disp.seeder, cell.K=K){
	lth<-length(pop.grid[1,1,1,1,])
	if (lth==1){
		print("Warning: Space consists of one grid cell - dispersal is meaningless")
		return(pop.grid)	
	}
	#get populations of stolons - Graeme, pls check that these are in the second row
	pop.size.grid<-pop.grid[2,1,1,,]
	ad.or<-apply(pop.grid[3:5,1,1,,], 2:3, sum)
	# work out number of dispersing stolons
	n.disp.grid<-ad.or/cell.K*(4*(lth-1))/lth^2 #dispersal probability scales with popsize
	n.disp.grid<-floor(n.disp.grid*pop.size.grid)
	# make arrays for the seeders
	minus.os<-array(0, dim=rep(lth, 2))
	minus.fr<-array(0, dim=rep(lth, 2))
	pl.os<-0
	pl.fr<-0
	# disperse the stolons and remove dispersing seeds from grid cells
	for (x in 1:lth){
		for (y in 1:lth){
			# obligate resprouters, nearest neighbour dispersal
			if (n.disp.grid[x, y]!=0){
				rows<-rmultinom(1, n.disp.grid[x,y], rep(1, 9))
				pop.grid[2,1,1,x,y]<-pop.grid[2,1,1,x,y]-sum(rows)
				for (i in 1:9){
					pop.grid[2,1,1,neighbour.grid[i,1,x,y],neighbour.grid[i,2,x,y]]<-rows[i]+pop.grid[2,1,1,neighbour.grid[i,1,x,y],neighbour.grid[i,2,x,y]]
				}
			}
			# seeders, global spread
			minus.os[x,y]<-floor(p.disp.seeder*pop.grid[1,1,2,x,y])
			pl.os<-pl.os+minus.os[x,y]
			minus.fr[x,y]<-floor(p.disp.seeder*pop.grid[1,1,3,x,y])
			pl.fr<-pl.fr+minus.fr[x,y]	
		}	
	}
	# disperse seeds
	rain.os<-array(rmultinom(1, pl.os, rep(1, lth*lth)), dim=c(lth, lth))
	rain.fr<-array(rmultinom(1, pl.fr, rep(1, lth*lth)), dim=c(lth, lth))
	pop.grid[1,1,2,,]<-pop.grid[1,1,2,,]-minus.os+rain.os
	pop.grid[1,1,3,,]<-pop.grid[1,1,3,,]-minus.fr+rain.fr
	# return pop.grid
	pop.grid
}

# takes fire grid and creates random fire events on the grid
# fires are of size cl.size*fire.scale so fire.scale must divide perfectly into space size
fire.grid<-function(f.grid, p.ignit, cond.p.fire, fire.scale, cell.size, sp.size){
	if (sp.size%%fire.scale!=0) {
		print("Warning: fire scale does not divide perfectly into space")
		return(NULL)	
	}
	lth.small<-sp.size/(cell.size*fire.scale)
	lth<-sp.size/cell.size
	small.grid<-f.grid[seq(1, lth, cell.size*fire.scale), seq(1, lth, cell.size*fire.scale)]
	small.grid<-small.grid+1
	small.grid[which(small.grid==2)]<-cond.p.fire
	small.grid<-rbinom(lth.small^2, 1, small.grid*p.ignit)
	big.grid<-rep(small.grid, each=fire.scale)
	big.grid<-as.vector(t(array(big.grid, dim=c(fire.scale*lth.small, lth.small))))
	big.grid<-rep(big.grid, each=fire.scale)
	big.grid<-array(big.grid, dim=rep(fire.scale*lth.small, 2))
	big.grid
}

prop.occ.collect<-function(pop.array, steps){
	if (steps>1){
		prop.occ<-c(sum(apply(grid[,1,1,,],2:3,sum)>0), sum(apply(grid[,1,2,,],2:3,sum)>0), sum(apply(grid[,1,3,,],2:3,sum)>0))
		prop.occ<-prop.occ/(steps^2)	
	}
	if (steps==1){
		prop.occ<-c(sum(grid[,1,1,,])>0, sum(grid[,1,2,,])>0, sum(grid[,1,3,,])>0)
	}	
	prop.occ	
}

